package org.apache.commons.graph.view.finder

import com.gitee.wsl.ext.map.putIfAbsent
import org.apache.commons.graph.bean.Edge
import org.apache.commons.graph.view.IDirectedGraphView


object SpanningTreeFinder {
    fun <V, E : Edge<V>> find(graph: IDirectedGraphView<V, E>): SpanningTree<V, E> {
        return findForEdges<V, E>(graph.edgeSet())
    }

    fun <V, E : Edge<V>> findForEdges(edges: Iterable<E>): SpanningTree<V, E> {
        val uf= UnionFind<V>()

        val treeEdges = ArrayList<E>()

        for (edge in edges) {
            val source: V = edge.source
            val target: V = edge.target

            if (uf.addVertex(source)) {
                uf.addVertex(target)
                uf.union(source, target)
                treeEdges.add(edge)
                continue
            }

            if (uf.addVertex(target)) {
                uf.union(source, target)
                treeEdges.add(edge)
                continue
            }

            val parent1 = uf.find(source)
            val parent2 = uf.find(target)
            if (parent1 == parent2) continue

            uf.union(source, target)
            treeEdges.add(edge)
        }
        return SpanningTree<V, E>(treeEdges)
    }
}

class SpanningTree<V, E : Edge<V>>(private val edges: List<E>) {
    val parentMap: Map<V, V>
        get() {
            val map = HashMap<V, V>(edges.size)
            for (edge in edges) {
                map[edge.target] = edge.source
            }
            return map
        }

    val roots: Set<V>
        get() {
            val set = LinkedHashSet<V>()
            for (edge in edges) {
                set.add(edge.source)
            }

            for (edge in edges) {
                set.remove(edge.target)
            }
            return set
        }

    val childrenMap: Map<V, Set<V>>
        get() {
            val map = HashMap<V, MutableSet<V>>()
            for (edge in edges) {
                val source: V = edge.source
                var set = map[source]
                if (set == null) {
                    set = LinkedHashSet()
                    map[source] = set
                }
                set.add(edge.target)
            }
            return map
        }
}


/**
 * 并查集是一种树型的数据结构，用于处理一些不交集（Disjoint Sets）的合并及查询问题。
 *
 *
 * Find：确定元素属于哪一个子集。它可以被用来确定两个元素是否属于同一子集。 Union：将两个子集合并成同一个集合。
 */
class UnionFind<V> {
    private val parentMap: MutableMap<V, V>
    private val rankMap: HashMap<V, Int>

    // 不相交的集合的个数
    var setCount: Int = 0
        private set

    constructor() {
        this.parentMap = LinkedHashMap()
        this.rankMap = HashMap()
    }

    constructor(vertices: Collection<V>) {
        this.parentMap = LinkedHashMap(vertices.size)
        this.rankMap = HashMap(vertices.size)
        for (vertex in vertices) {
            addVertex(vertex)
        }
    }

    fun addVertex(vertex: V): Boolean {
        if (parentMap.putIfAbsent(vertex, vertex) == null) {
            rankMap[vertex] = 0
            setCount++
            return true
        }
        return false
    }

    /**
     * 找到局部树的根节点
     */
    fun find(vertex: V): V? {
        var parent = parentMap[vertex]
        if (parent == null) return null

        if (parent == vertex) return parent

        var current: V? = parent
        do {
            parent = parentMap[current]
            if (parent == current) {
                break
            }
            current = parent
        } while (true)

        val root = current

        // 路径压缩，把路径上的所有节点都和根节点建立直接关联
        current = vertex
        while (current != root) {
            parent = parentMap.get(current)
            parentMap.put(current!!, root!!)
            current = parent
        }

        return root
    }

    fun union(v1: V, v2: V) {
        val parent1 = find(v1)
        val parent2 = find(v2)

        if (parent1 == null || parent2 == null) return

        if (parent1 == parent2) {
            return
        }

        val rank1: Int = rankMap[parent1]!!
        val rank2: Int = rankMap[parent2]!!
        if (rank1 > rank2) {
            parentMap.put(parent2, parent1)
        } else if (rank1 < rank2) {
            parentMap.put(parent1, parent2)
        } else {
            parentMap.put(parent2, parent1)
            rankMap.put(parent1, rank1 + 1)
        }

        setCount--
    }
}